from flask import Flask, jsonify, request, session
from flask_admin import Admin, AdminIndexView, expose
from flask_admin.contrib.sqla import ModelView
from flask_babel import Babel
from sqlalchemy import create_engine, func, case
from sqlalchemy.orm import sessionmaker, scoped_session
from sqlalchemy.ext.automap import automap_base
from sshtunnel import SSHTunnelForwarder
from jose import JWTError, jwt
from config import isLocalDB, table_name_mapping, column_name_mapping
from config import ALGORITHM, SECRET_KEY, ADMIN_CREDENTIALS


# 创建 Flask 应用
flask_app = Flask(__name__)

# 设置中文语言
flask_app.config['BABEL_DEFAULT_LOCALE'] = 'zh_CN'
babel = Babel(flask_app)

# 设置密钥
flask_app.secret_key = 'adminKey'

# 创建数据库引擎和会话
def create_db_engine():
    if isLocalDB:
        # 本地数据库连接
        db_url = "mysql+pymysql://ucas_user:ucas_password@localhost/ucas_food?charset=utf8"
        engine = create_engine(db_url)
        return engine, None
    else:
        # SSH 隧道连接
        ssh_tunnel = SSHTunnelForwarder(
            ("124.70.73.194", 22),  # SSH 服务器地址和端口
            ssh_username="root",  # SSH 用户名
            ssh_password="!@#123qwe",  # SSH 密码
            remote_bind_address=("localhost", 3306),  # 远程数据库地址和端口
        )
        ssh_tunnel.start()
        print("SSH tunnel started successfully.")

        # 数据库连接 URL
        db_url = f"mysql+pymysql://ucas_food:123456@127.0.0.1:{ssh_tunnel.local_bind_port}/ucas_food?charset=utf8"
        engine = create_engine(db_url)
        return engine, ssh_tunnel

# 创建数据库引擎和会话
engine, ssh_tunnel = create_db_engine()
Session = scoped_session(sessionmaker(bind=engine))

# Token 验证方法
def verify_token(token: str):
    credentials_exception = Exception("Invalid token")
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        user_id: str = payload.get("sub")
        if user_id is None:
            raise credentials_exception
        return user_id
    except JWTError:
        raise credentials_exception


@flask_app.before_request
def before_request():
    if request.path.startswith("/admin/admin"):
        token = request.args.get("token")  # 从 URL 参数中获取 Token
        if not token:
            # 返回 JSON 错误响应
            return jsonify({
                "status": "error",
                "message": "Token is missing. Please login first."
            }), 401

        # 调用 verify_token 方法验证 Token
        try:
            user_id = verify_token(token)
            if user_id not in ADMIN_CREDENTIALS:  # 检查用户是否是管理员
                # 返回 JSON 错误响应
                return jsonify({
                    "status": "error",
                    "message": "Permission denied. Only admin users can access this page."
                }), 403

            # 将用户信息存储在会话中
            session["user"] = user_id
        except Exception as e:
            # 返回 JSON 错误响应
            return jsonify({
                "status": "error",
                "message": "Invalid token. Please login again."
            }), 401


class MyAdminIndexView(AdminIndexView):
    @expose('/')
    def index(self):
        db_session = Session()

        # 用户总数
        user_count = db_session.query(Base.classes['user']).count()

        # 用户 BMI 分布
        bmi_distribution = db_session.query(
            case(
                (func.round(Base.classes['message'].weight / func.pow(Base.classes['message'].height / 100, 2), 1) < 18.5, '偏瘦'),
                (func.round(Base.classes['message'].weight / func.pow(Base.classes['message'].height / 100, 2), 1) < 24, '正常'),
                (func.round(Base.classes['message'].weight / func.pow(Base.classes['message'].height / 100, 2), 1) < 28, '超重'),
                (func.round(Base.classes['message'].weight / func.pow(Base.classes['message'].height / 100, 2), 1) >= 28, '肥胖'),
                else_='未知'
            ).label('bmi_category'),
            func.count(Base.classes['message'].id)
        )\
        .group_by('bmi_category')\
        .all()

        # 用户性别分布
        gender_distribution = db_session.query(
            Base.classes['message'].gender,
            func.count(Base.classes['message'].id)
        )\
        .group_by(Base.classes['message'].gender)\
        .all()

        # 用户年龄分布
        age_distribution = db_session.query(
            case(
                (Base.classes['message'].age < 18, '未成年'),
                (Base.classes['message'].age < 30, '18-29岁'),
                (Base.classes['message'].age < 40, '30-39岁'),
                (Base.classes['message'].age < 50, '40-49岁'),
                (Base.classes['message'].age >= 50, '50岁及以上'),
                else_='未知'
            ).label('age_group'),
            func.count(Base.classes['message'].id)
        )\
        .group_by('age_group')\
        .all()

        db_session.close()

        # 渲染模板并传递统计信息
        return self.render('admin/index.html',
                          user_count=user_count,
                          bmi_distribution=bmi_distribution,
                          gender_distribution=gender_distribution,
                          age_distribution=age_distribution)
# 初始化 Flask-Admin
admin = Admin(flask_app, name='果壳悦食管理平台', template_mode='bootstrap4', index_view=MyAdminIndexView())

# 自动反射数据库表
Base = automap_base()
Base.prepare(autoload_with=engine)

# 视图设置
class MyModelView(ModelView):
    column_display_pk = True  # 显式显示主键列

    def __init__(self, model, session, **kwargs):
        # 获取模型的所有列
        columns = [column.key for column in model.__table__.columns]

        # 设置列名映射
        if model.__table__.name in column_name_mapping:
            self.column_labels = {column: column_name_mapping[model.__table__.name].get(column, column) for column in columns}
            # print(f"Column labels: {self.column_labels}")
        # 设置所有列可搜索
        self.column_searchable_list = columns
        # 设置所有列可排序
        self.column_sortable_list = columns
        # 设置所有列可筛选
        self.column_filters = columns
        self.can_create = False  # 禁止新增数据
        self.can_edit = False    # 禁止编辑数据

        # 允许编辑字符串主键
        if model.__table__.name in ['user', 'message']: 
            self.form_widget_args = {
                'id': {
                    'readonly': False  # 允许编辑主键字段
                }
            }
        super(MyModelView, self).__init__(model, session, **kwargs)
    
# 将自动生成的模型类添加到 Flask-Admin
for table_name, model_class in Base.classes.items():
    viewName = table_name
    if table_name in table_name_mapping:
            viewName = table_name_mapping[table_name]
    admin.add_view(MyModelView(model_class, Session, name=viewName))

# 关闭数据库连接和 SSH 隧道
@flask_app.teardown_appcontext
def shutdown_session(exception=None):
    Session.remove()
    if ssh_tunnel is not None:
        ssh_tunnel.stop()
        print("SSH tunnel stopped.")
